"""
Angr-based Reverse CTF Framework
功能：
  1. 加载二进制文件
  2. 内存 patch（写入指令）
  3. 栈变量符号化
  4. 从指定地址开始执行
  5. Inline Hook 指定函数
"""

import angr
import claripy
from capstone import Cs, CS_ARCH_X86, CS_MODE_64

want_addr  =0
class ReverseCTFFramework:
    def __init__(self, binary_path, base_addr=0x400000):
        self.binary_path = binary_path
        self.base_addr = base_addr
        self.project = angr.Project(binary_path, load_options={'main_opts': {'custom_base_addr': base_addr}})
        self.state = None
        self.simgr = None
        self.hooks = {}
        # self.strlen = 33

    def load_binary(self):
        """加载二进制文件"""
        print(f"[+] 加载二进制文件: {self.binary_path}")
        print(f"[+] 架构: {self.project.arch.name}")
        print(f"[+] 入口点: 0x{self.project.entry:x}")

    def patch_instruction(self, addr, asm_bytes):
        """
        在内存中 patch 指令（例如 nop out check）
        :param addr: 要 patch 的虚拟地址（相对于基址）
        :param asm_bytes: 字节码，如 b"\x90\x90" (nop nop)
        """
        self.project.loader.memory.store(addr, asm_bytes)
        print(f"[Patch] 已在 0x{addr:x} 写入 {len(asm_bytes)} 字节")

    def patch_assemble(self, addr, asm_instructions, arch="x64"):
        """
        使用汇编指令 patch（需 capstone + keystone）
        示例: patch_assemble(0x401000, ["mov eax, 1", "ret"])
        """
        try:
            import keystone
            ks_arch, ks_mode = {
                "x64": (keystone.KS_ARCH_X86, keystone.KS_MODE_64),
            }[arch]
            ks = keystone.Ks(ks_arch, ks_mode)
            encoding, _ = ks.asm(asm_instructions)
            self.patch_instruction(addr, bytes(encoding))
        except ImportError:
            print("[-] 需要安装 keystone-engine: pip install keystone-engine")

    def setup_symbolic_stack_var(self, start_addr, size=8, name="sym_var"):
        """
        将栈上某变量符号化
        :param start_addr: 栈地址（如 state.regs.rbp - 0x10）
        :param size: 变量大小（字节）
        :param name: 符号名称
        """
        sym_val = claripy.BVS(name, size * 8)  # 8 bits per byte
        self.state.memory.store(start_addr, sym_val)
        print(f"[Symbolic] 栈地址 0x{start_addr:x} 已符号化为 {name}")
        return sym_val

    def start_from_address(self, addr):
        """从指定地址开始执行"""
        self.state = self.project.factory.blank_state(
            addr=addr,
            mode='symbolic',
            add_options={
                angr.options.LAZY_SOLVES,
                angr.options.SYMBOLIC_MEMORY_OPTIMIZATIONS
            }
        )
        self.simgr = self.project.factory.simulation_manager(self.state)
        print(f"[Exec] 从地址 0x{addr:x} 开始执行")

    def inline_hook(self, target_addr, hook_func, length=5):
        """
        内联 Hook：在目标函数开头插入跳转
        :param target_addr: 要 hook 的函数地址
        :param hook_func: Python hook 函数，接受 state 参数
        :param length: 覆盖的字节数（通常 5 字节用于 jmp）
        """
        def _hook(state):
            print(f"[Hook] 拦截执行: 0x{target_addr:x}")
            hook_func(state)

        # 保存原始指令（可选）
        orig_bytes = self.project.loader.memory.load(target_addr, length)
        self.hooks[target_addr] = orig_bytes

        # 插入 hook
        self.project.hook(target_addr, _hook, length=length)
        print(f"[Hook] 已在 0x{target_addr:x} 设置内联 hook")

    def run(self, step_limit=100):
        """运行仿真"""
        # if not self.simgr:
        #     raise Exception("请先调用 start_from_address()")
        # self.simgr.run(step_limit)
        # print(f"[Run] 执行完成，共 {len(self.simgr.active)} 个活跃状态")


        input_len = 24  # 假设 flag 不超过 16 字节
        flag_chars = [claripy.BVS(f'char_{i}', 8) for i in range(input_len)]
        flag = claripy.Concat(*flag_chars)  # 拼接成一个 bitvector

        project  = self.project
        # 创建初始状态
        # main_addr = project.loader.main_object.get_symbol('main').rebased_addr
        # state = project.factory.call_state(main_addr)
        state = project.factory.full_init_state(
            
        )
        # flag_data = project.loader.memory.load(0x400000+0x2120, 24) 
        
        # print("Flag:", flag_data.decode('ascii'))
        # state =project.factory.call_state('main')
        orginal_data = project.loader.memory.load(0x400000+0x182b,0x40)
        for i in range(0x40):  # 添加符号化输入
            print(f'{orginal_data[i]:x}',end=' ')
        print()
        ans = 'Bl@sting_1s_a_g00d_Way!!'
        # state.solver.add(flag.get_byte(0) == ord('B'))
        # state.solver.add(flag.get_byte(1) == ord('l'))
        # state.solver.add(flag.get_byte(2) == ord('@'))
        # state.solver.add(flag.get_byte(3) == ord('s'))
        # state.solver.add(flag.get_byte(4) == ord('t'))
        # state.solver.add(flag.get_byte(5) == ord('i'))
        # state.solver.add(flag.get_byte(3) == ord('g'))
        # state.solver.add(flag.get_byte(4) == ord('{'))
        # # # 最后一个为 '}'
        # state.solver.add(flag.get_byte(23) == ord('}'))
        for i in range(24):
            byte = flag.get_byte(i)  # 获取第 i 个字节
            # 约束：0x20 <= byte <= 0x7e
            if i > 2:
                is_lower  = claripy.And(byte >= 0x61, byte <= 0x7a)
                is_upper  = claripy.And(byte >= 0x41, byte <= 0x5a)
                is_digit  = claripy.And(byte >= 0x21, byte <= 0x39)
                is_underscore = (byte == 0x5f)
                # state.solver.add(byte !=0 )
                state.solver.add(is_lower | is_upper | is_digit | is_underscore )
            else:
                state.solver.add(byte  == ans[i])
        # === 关键：模拟 gets 输入 ===
        # 分配一块内存作为 gets 的目标缓冲区
        # buffer_addr = state.heap.allocate(64)
        # state.memory.store(buffer_addr, flag)  # 将符号化输入写入缓冲区

        # Hook gets 函数，使其不真正读取，而是使用我们的符号输入
        @project.hook(0x400000+0x1676)
        def gets_hook(state):
            # 将符号输入写入 gets 的参数（rdi 是第一个参数）
            target_addr = state.regs.rdi  # gets(buffer) 中 buffer 的地址
            state.memory.store(target_addr, flag)
            # 设置返回值为 buffer 地址
            state.regs.rax = target_addr
            print("是否有符号性?", state.memory.load(state.regs.rdi, 1).symbolic)  # 应为 True
            # 跳过真实 gets 执行
            # state.project.hooked_by(state.addr).clear()  # 防止递归
        
        def strlen_hook(state):
        # @project.hook(project.loader.find_symbol('strlen'))
        # def strlen_hook(state):
        #     # 将符号输入写入 gets 的参数（rdi 是第一个参数）
            # target_addr = state.regs.rdi  # gets(buffer) 中 buffer 的地址
            # state.memory.store(target_addr, flag)
            # 设置返回值为 buffer 地址
            #global strlen  
            state.regs.rax = 24
            print("strlen:")
            # print(f'rip {state.solver.eval(state.regs.rip):#x}')
            self.strlen = self.strlen -1
        def my_compare(state):
            
            rdx = state.memory.load(state.regs.rbp-0xa8, 8  )
            global want_addr
            want_addr = rdx
            mem_str = state.memory.load(rdx, 24)
            print('test',mem_str.symbolic)
            state.solver.add(mem_str == b'\x97\xd5\x60\x43\xb4\x10\x43\x73\x0f\xda\x43\xcd\xd3\xe8\x73\x4a\x94\xc3\xcd\x71\xbd\xdc\x97\x1a')
            #
        project.hook(0x400000+0x1891,my_compare,0x18e5-0x1891)
        class SymbolicStrlen(angr.SimProcedure):
            def __init__(self):
                super(SymbolicStrlen, self).__init__()
                self.length = claripy.BVS('strlen_result', 64)  # 64位系统
                
                # self.state.solver.add(self.length <= 16)
            def run(self, string):
                # 获取字符串地址
                addr = string

                # 创建一个符号变量表示长度
                

                # 可选：添加约束，比如长度在 8~16 之间
                self.state.solver.add(self.length <7)
                self.state.solver.add(self.length >2)
                # 但我们也要确保这个长度与实际字符串一致
                # angr 的 SimProcedure 已经处理了这个逻辑
                # 如果你强制返回符号变量而不验证，可能导致不一致

                # ⚠️ 注意：这里直接返回符号变量，**不检查实际字符串**
                # 仅用于特殊场景（如绕过长度检查）
                return 24

# Hook
        symbol_strlen = SymbolicStrlen()
        # project.hook_symbol('strlen', strlen_hook)
            #strlen = strlen - 1
        project.hook(0x400000+0x16A3,strlen_hook)
        project.hook(0x400000+0x14AE,strlen_hook)
        # project.hook_symbol('strlen',strlen_hook)

            # print("是否有符号性?", state.memory.load(state.regs.rdi, 1).symbolic)  # 应为 True
        # 创建 simulation manager
        simgr = project.factory.simulation_manager(state)

        # 查找输出 "right" 的路径
        def found_right(state):
            try:
                # 检查标准输出是否包含 "right"
                stdout_output = state.posix.dumps(1)
                return b"right" in stdout_output
            except:
                return False

        # 执行符号执行
        simgr.explore(find=0x400000+0x18ee,avoid=lambda s: b"error" in s.posix.dumps(1))

        # 输出结果
        if len(simgr.found) > 0:
            found_state = simgr.found[0]
            print("[+] 找到输出 'right' 的路径！")

            # 求解输入
            try:
                # print(f"[+] 输入为: {flag}")
                # sint = found_state.solver.eval(symbol_strlen.length,cast_to=int)
                # print(f"[+] 输入为: {sint}")
                str111 = found_state.memory.load(want_addr,24)
                str111aaa = found_state.solver.eval(str111,cast_to=bytes)
                print(str111aaa)
                solutions = found_state.solver.eval_upto(flag, n=1000,cast_to=bytes)
                # 去除填充的 null 字节
                # print(f"[+] 输入为: {solution}")
                for solution in solutions:
                    # if solution.endswith(b'\x00'):
                    #     break
                    print(f"[+] 输入为: {solution}")
                    # solution = solution.rstrip(b'\x00').decode()
                    # print(f"[+] 输入为: {solution}")
            except Exception as e:
                print(f"[-] 无法求解输入: {e}")
        else:
            print("[-] 未找到输出 'right' 的路径")
    def run_1(self, step_limit=100):
        """运行仿真"""



        

  
        project  = self.project

        list1 = [0]*256
        findal_res = [0x97 ,0xd5 ,0x60 ,0x43 ,0xb4 ,0x10 ,0x43 ,0x73 ,0x0f ,0xda ,0x43 ,0xcd ,0xd3 ,0xe8 ,0x73 ,0x4a ,0x94 ,0xc3 ,0xcd ,0x71 ,0xbd ,0xdc ,0x97 ,0x1a]
        for i in range(256):
            func_addr = 0x4013E1
            state = project.factory.call_state(func_addr, i ,ret_addr=0x100000  )
            simgr = project.factory.simulation_manager(state)
            simgr.run()
            # simgr.run()
            print(simgr)
            for err in simgr.errored:
                print("错误原因:", err.error)
                # print("最后地址:", hex(err.addr))
                for s in [err.state]:
                    angr_v = s.solver.eval(s.regs.rax) & 0xff
                    if angr_v in findal_res:
           
                        list1[angr_v] = i

        for i in findal_res:
            if list1[i] != 0:
                
                print(f"{list1[i]},",end='')
    def run_2(self, step_limit=100):
        """运行仿真"""
        project  = self.project

        res_1=[132,177,2,55,142,76,55,206,125,137,55,235,88,175,206,192,129,70,235,213,176,242,132,9]
        
        res_str = ""
        res_str2 =''
        for i in range(24):
            for j in range(0x20,0x7e):
                func_addr = 0x4012a9
                state = project.factory.call_state(func_addr, j  ,(i%7)+1,ret_addr=0x100000  )
                simgr = project.factory.simulation_manager(state)
                simgr.run()
                # simgr.run()
                print(simgr)
                for err in simgr.errored:
                    print("错误原因:", err.error)
                    # print("最后地址:", hex(err.addr))
                    for s in [err.state]:
                        angr_v = s.solver.eval(s.regs.rax) & 0xff
                        if angr_v == res_1[i]:
         
                            res_str2 += str(i)
                            res_str += chr(j)
        print(res_str)
        print(res_str2)



    def dump_state(self):
        """打印当前状态信息"""
        if self.state:
            print("[State] 寄存器状态:")
            print(f"  RIP: 0x{self.state.addr:x}")
            print(f"  RAX: 0x{self.state.regs.rax.args[0]:x}")
            print(f"  RSP: 0x{self.state.regs.rsp.args[0]:x}")
            print("[Memory] 符号化变量:")
            for var in self.state.memory.variables:
                print(f"  {var}")


# ======================
# 使用示例
# ======================
def my_hook(state):
    print("[Custom Hook] 调用被拦截！可以修改寄存器或内存")
    state.regs.rax = 1  # 强制返回 1


if __name__ == "__main__":
    # 创建框架实例
    framework = ReverseCTFFramework("./hardtest")

    # 1. 加载二进制
    framework.load_binary()

    # 2. Patch 指令：例如 nop 掉一个跳转
    # framework.patch_instruction(0x400000+0x15e9, b"\x90"*(0x1660-0x15e9))  # nop nop
    # framework.patch_instruction(0x400000+0x17df, b"\x90"*(0x1801-0x17df))  # nop nop
    # framework.patch_instruction(0x400000+0x1711, b"\x90"*(0x173B-0x1711))  # nop nop

    # framework.patch_instruction(0x400000+0x15e9, b"\x90"*(0x1660-0x16f2))  # nop nop
    # framework.patch_instruction(0x400000+0x16f2, b"\x90"*(0x1765-0x16f2)) 
    # framework.patch_instruction(0x400000+0x178c, b"\x90"*(0x182B-0x178c)) 

    # # 或使用汇编 patch（需 keystone）
    # # framework.patch_assemble(0x401130, ["mov eax, 0", "jmp 0x402000"])

    # # 3. 从某个地址开始执行（如 bypass anti-debug）
    # framework.start_from_address(0x4011c0+0x1531)

    # # 4. 将栈上变量符号化（假设在 rbp-0x20）
    # rbp = framework.state.regs.rbp
    # var_addr = rbp - 0x20
    # sym_input = framework.setup_symbolic_stack_var(var_addr, size=16, name="user_input")

    # # 5. 设置 inline hook
    # framework.inline_hook(0x403000, my_hook, length=5)  # hook check_func

    # # 6. 运行
    framework.run_2(step_limit=200)
    
    # # 7. 查看状态
    # framework.dump_state()
    # for j in range(0x20,0x7e):
    #     print(chr(j))
    # 可继续使用 solver 求解输入
    # solution = framework.state.solver.eval(sym_input, cast_to=bytes)
'''
Bgg0y!
0 7 14 15 21 22



Bl@st1ng_1s_a_g00d_Way!!
'''